#include "suffixTree.hpp"

int sNode::_count = 0;
int numFromLetter (char ch)
{
    if (ch >= 97 && ch <= 122)
	ch -= 32;
    return ch - 65 + 1;
}

int hashCode(int num, char ch)
{
    return num*100 + numFromLetter(ch);
}

int hashCode(int num, int ch)
{
    return num*100 + ch;
}

int findPatMatch(const std::string &str1, const std::string &str2)
{
    int i = 0;
    while (i < str1.size()) {
	if (str1[i] == str2[i])
	    ++i;
	else
	    break;
    }
    return i;
}

void printRange (std::ofstream &out, std::string &str, int lb, int ub)
{
    for (int i = lb; i <= ub; i++) {
	if (str[i] == '\0')
	    break;
	out<<str[i];
    }
}

void sNode::serialize (std::ofstream &out, _transMap *map, std::vector<std::string*> &words)
{
    /*find the edges going out from the root in alphabetical order*/
    out<<"(";
    for (int i = 1; i <= 26; i++) {
	_transMap::iterator itr = map->find(hashCode(_id, i));
	if (itr == map->end())
	    continue;
	sEdge *e = itr->second;
	out<<"{";
	printRange(out, *(words[e->wordIndex()]), e->startIndex(), e->endIndex());
	out<<"}";
	e->to()->serialize(out, map, words);
    }
    out<<")";

    return;
}

std::pair<sNode*, int> sNode::walkTree(std::string &word, int pntr, _transMap
	&transMap, std::vector<std::string*> &words)
{
    int idx = hashCode(_id, word[pntr]);
    if (transMap.find(idx) == transMap.end())
	return std::pair<sNode*, int>(this, pntr);
    sEdge &e = *(transMap[idx]);
    std::string &tmp = *(words[e.wordIndex()]);
    int j = pntr;
    bool flag = false;
    for (int i = e.startIndex(); i <= e.endIndex(); i++) {
	if (tmp[i] != word[j]) {
	    flag = true;
	    break;
	}
	++j;
    }
    if (flag)
	return (j > pntr) ? std::pair<sNode*, int>(e.to(), j) : std::pair<sNode*,
	       int>(this, j);
    return e.to()->walkTree(word, pntr+tmp.size(), transMap, words);
}

std::pair<sNode*, int>* suffixTree::canonize(sNode *s, int k, int p, int idx)
{
    if (p < k)
	return new std::pair<sNode*, int>(s, k);
    std::string &str = *((*words)[idx]);
    sEdge *leavEdge = (*transMap)[hashCode(s->id(), str[k])];
    while ((leavEdge->endIndex() - leavEdge->startIndex()) <= p-k) {
	k += leavEdge->endIndex() - leavEdge->startIndex() + 1;
	s = leavEdge->to();
	if (k <= p)
	    leavEdge = (*transMap)[hashCode(s->id(), str[k])];
    }
    return new std::pair<sNode*, int>(s, k);
}

std::pair<bool, sNode*>* suffixTree::testAndSplit(sNode *s, int k, int p, char t, int idx)
{
    if (k <= p) {
	std::string &str = *((*words)[idx]);
	int key1 = hashCode(s->id(), str[k]);
	sEdge *leavEdge = (*transMap)[key1];
	if (t == str[leavEdge->startIndex()+p-k+1])
	    return new std::pair<bool, sNode*>(true, s);
	else {
	    sEdge *e = new sEdge(s, leavEdge->startIndex(), leavEdge->startIndex()+p-k, idx);
	    sNode *r = new sNode(e);
	    e->to(r);
	    (*transMap)[key1] = e;
	    e = new sEdge(r, leavEdge->startIndex()+p-k+1, leavEdge->endIndex(), idx);
	    e->to(leavEdge->to());
	    (*transMap)[hashCode(r->id(), str[leavEdge->startIndex()+p-k+1])] = e;
	    return new std::pair<bool, sNode*>(false, r);
	}
    } else if ((*transMap)[hashCode(s->id(), t)])
	return new std::pair<bool, sNode*>(true, s);
    else
	return new std::pair<bool, sNode*>(false, s);
}

std::pair<sNode*, int>* suffixTree::update(sNode *s, int k, int i, int idx)
{
    sNode *oldr = _root;
    std::string &str = *((*words)[idx]);
    std::pair<bool, sNode*> *testResult = testAndSplit(s, k, i - 1, str[i], idx);
    while (!testResult->first) {
	sEdge *leavEdge = new sEdge(testResult->second, i, str.size(), idx);
	sNode *destNode = new sNode(leavEdge);
	leavEdge->to(destNode);
	(*transMap)[hashCode(testResult->second->id(), str[i])] = leavEdge;
	if (oldr != _root)
	    oldr->sLink(testResult->second);
	oldr = testResult->second;
	std::pair<sNode*, int>* canonized = canonize(s->sLink(), k, i-1, idx);
	s = canonized->first;
	k = canonized->second;
	testResult = testAndSplit(s, k, i-1, str[i], idx);
    }
    if (oldr != _root)
	oldr->sLink(s);

    return new std::pair<sNode*, int>(s, k);
}

void suffixTree::insert(std::string &toInsert)
{
    std::string *_str = new std::string(toInsert);
    std::string &str = *_str;
    words->push_back(&str);
    int idx = words->size() - 1;
    /*add transitions from _exp to root for all the chars in str*/
    for (int i = 0; i < str.size(); i++) {
	int key = hashCode(_exp->id(), str[i]);
	if ((*transMap)[key])
	    continue;
	sEdge *e = new sEdge(_exp, _root, i, i);
	(*transMap)[key] = e;
    }
    sNode *s = _root;
    int k = 0;
    int startIdx = 0;
    if (idx > 0) { /*set the starting index for the idx>=1 word being inserted
		     to the suffix tree*/
	std::pair<sNode*, int> tmp = s->walkTree(str, 0, *transMap, *words);
	s = tmp.first;
	k = startIdx = tmp.second;
    }
    for (int i = startIdx; i < str.size(); i++) {
	std::pair<sNode*, int> *updateRes = update(s, k, i, idx);
	std::pair<sNode*, int> *canoRes = canonize(updateRes->first, updateRes->second, i, idx);
	s = canoRes->first;
	k = canoRes->second;
    }

    return;
}

void suffixTree::serialize (std::ofstream &out)
{
    _root->serialize(out, transMap, *words);
    out<<std::endl;
    return;
}

